package com.heima.algorithm.sort;

import java.util.concurrent.ThreadLocalRandom;

/**
 * 快速选择算法，在数组元素中去选第排名为k的的元素(排名从0开始)
 * @author 勾新杰
 * @version 1.0
 * @date 2024/11/13 20:00
 */
public class QuickSelect {

    /**
     * 快速选择算法，核心思想是快速排序的分区思想
     * @param arr 指定数组
     * @param k 排名从0开始
     * @return 第k名的元素
     */
    public static int quickSelect(int[] arr, int k) {
        return quickSelect(arr, 0, arr.length - 1, k);
    }

    /**
     * 快速选择算法，核心思想是快速排序的分区思想O(n)
     * @param arr 指定数组
     * @param begin 开始下标（包含）
     * @param end 结束下标（包含）
     * @param k 排名从0开始
     * @return 第k名的元素
     */
    private static int quickSelect(int[] arr, int begin, int end, int k) {
        int pivot = partition(arr, begin, end);
        if (pivot == k){
            return arr[pivot];
        } else if (pivot < k) {
            return quickSelect(arr, pivot + 1, end, k);
        } else {
            return quickSelect(arr, begin, pivot - 1, k);
        }
    }

    /**
     * 分区，返回基准点的索引
     * @param arr 指定数组
     * @param begin 开始下标（包含）
     * @param end 结束下标（包含）
     * @return 基准点的索引
     */
    private static int partition(int[] arr, int begin, int end) {
        int randomIndex = ThreadLocalRandom.current().nextInt(end - begin + 1) + begin;
        swap(arr, begin, randomIndex);
        int pivot = arr[begin];
        int le = end; // le指针找到小于等于基准点的元素的位置，找到后还应该停下来，后面把这个元素放到基准点前面
        int ge = begin + 1; // ge指针找到大于等于基准点的元素的位置，找到后还应该停下来，后面把这个元素放到基准点后面
        // 所以上述方案最终基准点两边都有可能有和基准点元素相等的
        while (ge <= le){ // ge <= le，因为ge和le相等的时候是有可能找到的就是相等的情况
            // 找到比基准点小或等的元素，找到后应该停下来，后面把这个元素放到基准点前面
            while (ge <= le && arr[le] > pivot){ // 不能交换
                le--;
            }
            // 找到比基准点大或等的元素，找到后应该停下来，后面把这个元素放到基准点后面
            while (ge <= le && arr[ge] < pivot){ // 不能交换
                ge++;
            }
            // 还处于循环条件内，则交换两个元素并让ge++ le--
            if (ge <= le){
                if (arr[ge] != arr[le]) {
                    swap(arr, ge, le);
                }
                ge++;
                le--;
            }
        }
        swap(arr, begin, le); // 将基准点元素放到合适位置
        return le; // 返回基准点元素索引
    }

    /**
     * 交换两个元素
     * @param arr 指定数组
     * @param index1 索引1
     * @param index2 索引2
     */
    private static void swap(int[] arr, int index1, int index2) {
        int temp = arr[index1];
        arr[index1] = arr[index2];
        arr[index2] = temp;
    }

    /**
     * 寻找中位数
     * @param nums 指定数组
     * @return 中位数
     */
    public static int findMedian(int[] nums) {
        int half = nums.length >> 1;
        if ((nums.length & 1) == 1){
            return quickSelect(nums, half);
        }else {
            int left = quickSelect(nums, half);
            int right = quickSelect(nums, half - 1);
            return (left + right) >> 1;
        }
    }
}
